# -*- coding: utf-8 -*-

import torch
import segmentation_models_pytorch as smp

class Accuracy(torch.nn.Module):
    def __init__(self, threshold=0.5, activation=None, ignore_channels=None):
         super(Accuracy, self).__init__()
         self.accuracy = smp.utils.metrics.Accuracy(threshold, activation, ignore_channels)
         self.__name__ = "accuracy"
         
    def forward(self, y_pr, y_gt):
        y_gt = y_gt.to(y_pr.device)
        y_pr = torch.argmax(y_pr, dim=1)
        y_pr = torch.unsqueeze(y_pr, dim=1)
        return self.accuracy.forward(y_pr, y_gt)
    
class IoU(torch.nn.Module):
    def __init__(self):
        super(IoU, self).__init__()
        self.__name__ = "fwiou"
    
    def forward(self, y_pr, y_gt):
        y_gt = y_gt.to(y_pr.device)
        
        classes = int(y_pr.shape[1])
        # classes = 3
        y_pr = torch.argmax(y_pr, dim=1)
        y_pr = torch.unsqueeze(y_pr, dim=1)
        
        iou, tot_positive = 0, 0
        for cl in range(1, classes):
            inter = torch.sum(torch.logical_and((y_pr == cl), (y_gt == cl)))
            union = torch.sum(torch.logical_or((y_pr == cl), (y_gt == cl)))
            positive = torch.sum((y_gt == cl))
            tot_positive += positive
            iou += (positive * inter / union) if union != 0 else positive
           
        fwiou = iou / tot_positive if tot_positive != 0 else torch.tensor(1)
        return fwiou


if __name__ == "__main__":
    from CC_CCI_dataset import get_TCSM_dataloader
    
    IMAGE_DIR = "D:\data\ct_lesion_seg\image"
    MASK_DIR = "D:\data\ct_lesion_seg\mask"
    CLASSES = 3
    BATCH_SIZE = 5
    LABELED_RATE = 0.5
    times = int(1 / LABELED_RATE - 1)
    
    
    train_loader, valid_loader, unlabeled_loader = get_TCSM_dataloader(IMAGE_DIR, MASK_DIR, 
                                                                        labeled_batch_size=BATCH_SIZE,
                                                                        unlabeled_times=times,
                                                                        device='cuda:0')
    
    for i, item in enumerate(valid_loader):
        if i >= 1:
            break
        x, y = item
    
    iou = IoU()(y, y)
    accuracy = Accuracy()(y, y)
    
    